import os
import numpy as np
import tensorflow as tf
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
from tensorflow.keras.models import Model
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# 设置常量
IMAGE_SIZE = 224
BATCH_SIZE = 32
EPOCHS = 10

# 垃圾分类类别
CLASSES = ['可回收垃圾', '有害垃圾', '厨余垃圾', '其他垃圾']

def create_model():
    # 加载预训练的MobileNetV2模型
    base_model = MobileNetV2(weights='imagenet', include_top=False, input_shape=(IMAGE_SIZE, IMAGE_SIZE, 3))
    
    # 冻结基础模型的层
    for layer in base_model.layers:
        layer.trainable = False
    
    # 添加自定义层
    x = base_model.output
    x = GlobalAveragePooling2D()(x)
    x = Dense(1024, activation='relu')(x)
    predictions = Dense(len(CLASSES), activation='softmax')(x)
    
    # 构建完整模型
    model = Model(inputs=base_model.input, outputs=predictions)
    return model

def train_model(train_dir, validation_dir):
    # 数据增强
    train_datagen = ImageDataGenerator(
        rescale=1./255,
        rotation_range=20,
        width_shift_range=0.2,
        height_shift_range=0.2,
        horizontal_flip=True,
        fill_mode='nearest'
    )
    
    validation_datagen = ImageDataGenerator(rescale=1./255)
    
    # 创建数据生成器
    train_generator = train_datagen.flow_from_directory(
        train_dir,
        target_size=(IMAGE_SIZE, IMAGE_SIZE),
        batch_size=BATCH_SIZE,
        class_mode='categorical',
        classes=CLASSES
    )
    
    validation_generator = validation_datagen.flow_from_directory(
        validation_dir,
        target_size=(IMAGE_SIZE, IMAGE_SIZE),
        batch_size=BATCH_SIZE,
        class_mode='categorical',
        classes=CLASSES
    )
    
    # 创建模型
    model = create_model()
    
    # 编译模型
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
        loss='categorical_crossentropy',
        metrics=['accuracy']
    )
    
    # 训练模型
    history = model.fit(
        train_generator,
        steps_per_epoch=train_generator.samples // BATCH_SIZE,
        epochs=EPOCHS,
        validation_data=validation_generator,
        validation_steps=validation_generator.samples // BATCH_SIZE
    )
    
    # 保存模型
    model.save('garbage_classifier.h5')
    return history

if __name__ == '__main__':
    # 设置训练和验证数据目录
    train_dir = 'data/train'
    validation_dir = 'data/validation'
    
    # 训练模型
    history = train_model(train_dir, validation_dir)